©PaperWeekly 原创 · 作者 | 苏剑林
正如 “XXX is all you need” 一样，有不少论文都以“简单得令人尴尬”命名（An Embarrassingly Simple XXX），但在笔者看来，这些论文大多数都是噱头多于实力。不过，笔者最近阅读到的一篇论文，真的让人不由得发出“简单得令人尴尬”的感叹～
论文的标题是《Finite Scalar Quantization: VQ-VAE Made Simple》[1]，顾名思义，这是一篇旨在用 FSQ（Finite Scalar Quantization）简化 VQ-VAE 的工作。随着生成模型、多模态 LLM 的逐渐流行，VQ-VAE 及其后续工作也作为“图像的 Tokenizer” 而“水涨船高”。然而，VQ-VAE 的训练本身也存在一些问题，而 FSQ 这篇论文则声称通过更简单的“四舍五入”就可以达到同样的目的，并且有着效果更好、收敛更快、训练更稳的优点。FSQ 真有这么神奇？接下来我们一起学习一下。
首先，我们来了解一下 “VQ”。VQ 全称是 “Vector Quantize”，可以翻译为“向量量子化”或者“向量量化”，是指将无限、连续的编码向量映射为有限、离散的整数数字的一种技术。如果我们将 VQ 应用在自编码器的中间层，那么可以在压缩输入大小的同时，让编码结果成为一个离散的整数序列。假设自编码器的重构损失能够让我们满意，那么这个整数序列就是原始图像的等价物，所有关于原始图像的操作都可以转化为整数序列上的操作。比如我们想训练图像生成模型，就只需要训练整数序列生成模型，而这跟本文生成等价，所以我们可以用它来训练一个 GPT，模型和流程都跟文本一模一样，训练完成后，我们就可以从 GPT 模型中采样整数序列，然后送到解码器中得到图像，从而完完成了图像生成模型的构建。说白了，“VQ +自编码器”将任意输入都转化为跟文本一致的整数序列，统一了不同模态数据的输入形式，同时也统一了它们的处理和生成模型。而这样的一个带有 VQ 功能的自编码器，就被称为 “VQ-VAE”。
早在四年前的文章《VQ-VAE的简明介绍：量子化自编码器》[2] 中我们就介绍过了 VQ-VAE，尽管被冠以 “VAE（Variational AutoEncoder）”之名，但它实际上跟 VAE 没啥关系，如上一节所说，它只是一个带有 VQ 功能的 AE（AutoEncoder）。既然是 AE，那么有 encoder 和 decoder，一个普通的 AE 是这样的：
让我们来逐步解释一下。首先，第一步是相同的，输入 到 encoder 中，输出编码向量 。然而，我们并不是直接将 输入到 decoder 中，而是先维护一个编码向量表 （Codebook），从中选取与 最相近的一个 送入到 decoder 中进行重构 。由于编码表是有限的，所以我们也可以将实际的编码结果理解为一个整数（即与 最相近的 的 ），这就是 VQ-VAE 中 “VQ” 的含义。当然，实际应用中，为了保证重构的清晰度，encoder 的输出可能是多个向量，每个向量经历同样的量化步骤变成一个整数，所以结果就是一张原本在连续实数空间的图片，被编码 VQ-VAE 编码为了一个整数的序列，这跟文本 Tokenizer 的作用是类似的，所以就有了“图像的 Tokenizer” 的说法。
然而，由于整个前向计算的流程中出现了 ，所以梯度无法回传到 encoder，这意味着我们无法优化 encoder。此时常见的手段是 Gumbel Softmax，但 Gumbel Softmax [3] 的效果通常也是次优的，所以作者巧妙地借助了 Straight-Through 为 VQ-VAE 设计了更好的梯度。可以说，这是 VQ-VAE 最精彩的内容，它告诉我们什么 “Attention is all you need” 都是虚的，“Gradient” 才是真正的 “all we need”！具体来说，VQ-VAE 利用了深度学习框架基本上都自带的 stop_gradient（即公式中的 ）函数来自定义梯度，所有经过 的输入，都会保持同样的输出，但梯度被强迫为零。所以对于式（2）中的 ，我们有这样一来，送入 decoder 的还是量化过后的 ，但优化器求梯度时用的是 ，而 是 encoder 出来的，所以 encoder 也能够被优化了。这个操作就叫做 “Straight-Through Estimator（STE）”，是为神经网络的不可导模块设计梯度的常用技巧之一。由于基于梯度的优化器依然是当前的主流，所以直接设计梯度往往比设计 loss 更贴近本质，当然通常也更难、更让人由衷地赞叹。
不顾，事情还没完，此时有两个问题：1、现在 encoder 是有梯度了，但是编码表 却没了梯度；2、 虽然可以随意定义梯度，但不是胡乱定义一个梯度都可以成功优化模型的。从（3）可以看成，要使它在数学上严格成立，那么唯一的解是 ，这告诉我们如果 STE 是合理的，那么 与 至少是相近的。于是为了梯度的合理性，同时也为了优化编码表，我们还可以补充一项辅助 loss：这样既可以迫使 与 接近，又可以让 也拥有了梯度，一举两得！但细想之下，还是有点美中不足：理论上 encoder 和 decoder 的重构 loss 已经足够优化 了，所以额外引入的一项应该主要用来优化 ，而不应该反过来明显影响 。为此，我们再次利用 技巧，不难证明式（4）的梯度等价于第一项把 的梯度停掉了，剩下 的梯度，第二项则反过来，目前两项是 1:1 的权重求和，意味着两项相同程度地相互影响，而刚才我们说了，这辅助 loss 应该主要用来优化 而不是 ，所以我们引入 ，将辅助 loss 改为然后再加到重构 loss 中，就得到了 VQ-VAE 总的 loss 了除此之外， 的优化还有另外的方案：首先将式（6）的 置零，这样一来 就又没有梯度了；然后我们观察到，VQ-VAE 的 VQ 操作其实跟 K-Means 聚类是有点相似的， 相当于是 个聚类中心。根据我们对 K-Means 的了解，聚类中心等于该类的所有向量的平均，所以 的一种优化方案就是 的滑动平均这等价于指定使用 SGD 优化 这一项 loss（其他项可以用 Adam 等）。该方案被 VQ-VAE-2 [4] 所使用。
可能有些读者疑惑，本文的主题不是 FSQ 吗？前面介绍 VQ-VAE 的篇幅是不是有点多了？事实上，由于 FSQ 完全对得起“简单得令人尴尬”这个评价，相比 VQ-VAE，介绍 FSQ 只需要“寥寥几行”，所以 VQ-VAE 不写长点，这篇博客就没几个字了哈～ 当然，将 VQ-VAE 写详细一点，也能让大家更深刻体会到 FSQ 的简单。准确来说，FSQ 只是用来替代 VQ-VAE 中的 “VQ” 的，它的离散化思路非常非常简单，就是“四舍五入”。首先，假设我们有一个标量 ，我们定义：这里的 是一个超参数， 就是 sigmoid 函数（原论文用了 ，笔者认为用 sigmoid 更科学）， 就是四舍五入为一个整数，所以不难看出 ，即 FSQ 运算将输出限制在了 个整数之中，从而实现了离散化。当然，多数情况下一个标量还不够，对于 ，每一维可以执行 FSQ 运行，于是即 维向量 被离散为 个整数之一。但要注意， 操作同样是没有梯度的（或者说梯度为零），不过经过 VQ-VAE 的铺垫，有些读者可能已经猜到接下来要做什么了：同样是利用 STE 技巧即反向传播用 之前的 求梯度。由于 前后本身是数值近似的，所以 FSQ 不需要额外 loss 来迫使近似的出现，也没有额外的编码表需要更新，FSQ 的简洁可见一斑！
如果将 VQ 理解为直接将编码向量聚类为 个不同的类别，那么 FSQ 就是将编码向量归纳出 个属性，每个属性划分为了 个等级，从而直接表达了 个不同的整数。当然，从最一般的考虑，每个属性的等级数也可以是不相同的 ，从而不同的组合数为 。按照原论文的建议 （之前的 LFQ [5] 则相当于 ），所以如果要对齐 VQ-VAE 的编码数量 的话，对于 FSQ 来说应该有 ，即 FSQ 对编码向量的维度 是有限制的（一般就只是个位数），并且通常是远小于 VQ-VAE 的编码维度（一般是三位数），这个直接后果是当编码总数 比较小（从而 也比较小）时，FSQ 的效果通常不如 VQ：▲ 不同编码表大小下 VQ 与 FSQ 的效果差异从图上可以看到，当编码表大小在 1000 左右时，FSQ 与 VQ 的效果接近；当编码表大小明显超过 1000 时，FSQ 占优；反之当编码表大小明显小于 1000 时，则 VQ 占优，这跟笔者自己的实验结果相近。笔者的参考代码为：https://github.com/bojone/FSQ
其他实验就是比较常规的证明 FSQ 比 VQ 更优异的各种任务实验了，读者自行阅读原论文就好。
从形式上来看，假设 ，那么 VQ 就好比是“一个 类的分类器”，而 FSQ 则是“ 个 级的打分器”，不管是从参数量、几何直观或者表达能力来看，其实 FSQ 都不如 VQ，但为什么 FSQ 有机会取得比 VQ 更好的结果呢？笔者认为有两方面的原因。第一个原因，是 encoder 和 decoder 太强。虽然 FSQ 本身弱一些，但是 encoder 和 decoder 都足够强了，所以基于神经网络的万能拟合能力假设，FSQ 相对于 VQ 的劣势，完全可以在 encoder 和 decoder 中弥补过来。而在 的设定下，两者的离散化程度都是一样的，也就是说 encoder 与 decoder 之间的“信息瓶颈”是一样的，因此 FSQ 本身的问题就显得微不足道了。第二个原因，是 VQ 的“队友”（梯度）太弱。VQ 的经典问题是编码表坍缩：当编码表增大时，编码表并没有被充分利用起来，反而由于恶性竞争导致编码表聚集到一块了，经典表现就是一个 5000 的编码表，最终效果还不如 500 的编码表。归根结底，这是梯度不够合理所致，尽管 VQ 已经巧妙地设计了梯度，但对于 这种硬指派的运算，基于梯度的优化都存在“赢者通吃”问题，这是坍缩的根本原因，而 FSQ 的 运算并不涉及到指派，它是直接取的近似值。当然，往大了讲，其实跟 VQ 类似的 K-Means 经常也有聚类中心坍缩的问题，可见 难优化已经是一个老大难的问题了。因此与其说 FSQ 太强，倒不如说是 VQ 的“队友”太弱。从以上两点分析可以看出，FSQ 要想超过 VQ，除了编码表要足够大之外，还有 encoder 与 decoder 要足够复杂，但这并非总能满足，比如有些场景下，我们希望模型的每一层输出都被量化，这时候平摊下来的 encoder 和 decoder 未必足够复杂，此时 FSQ 本身的不足就成为效果的瓶颈了。此外，VQ 之后的向量维度没有变化，可以是任意多维，而 FSQ 之前的向量必须投影到 维，这是很严重的降维，当我们需要用到投影之前的高维度近似向量时，就很难靠 FSQ 之后的低维向量简单恢复过来。所以，如果单纯是作为“图像的 Tokenzier”，那么 FSQ 或许已经可以取代 VQ，但这并不意味着任意场景下 VQ 都可以被 FSQ 取代。
本文介绍了 VQ-VAE 的 “VQ” 的一个及其简单的替代品—— FSQ（Finite Scalar Quantization），它直接通过四舍五入来对连续向量进行离散化，并且不需要额外的 loss 进行辅助。实验结果表明，当编码表足够大时，FSQ 比 VQ 更有优势。
[1] https://arxiv.org/abs/2309.15505
[2] https://kexue.fm/archives/6760
[3] https://kexue.fm/archives/6705[4] https://arxiv.org/abs/1906.00446
[4] https://arxiv.org/abs/2310.05737
#投 稿 通 道#
让你的文字被更多人看到
如何才能让更多的优质内容以更短路径到达读者群体，缩短读者寻找优质内容的成本呢？答案就是：你不认识的人。
总有一些你不认识的人，知道你想知道的东西。PaperWeekly 或许可以成为一座桥梁，促使不同背景、不同方向的学者和学术灵感相互碰撞，迸发出更多的可能性。
PaperWeekly 鼓励高校实验室或个人，在我们的平台上分享各类优质内容，可以是最新论文解读，也可以是学术热点剖析、科研心得或竞赛经验讲解等。我们的目的只有一个，让知识真正流动起来。
📝 稿件基本要求：
• 文章确系个人原创作品，未曾在公开渠道发表，如为其他平台已发表或待发表的文章，请明确标注
• 稿件建议以 markdown 格式撰写，文中配图以附件形式发送，要求图片清晰，无版权问题
• PaperWeekly 尊重原作者署名权，并将为每篇被采纳的原创首发稿件，提供业内具有竞争力稿酬，具体依据文章阅读量和文章质量阶梯制结算
📬 投稿通道：
• 投稿邮箱：hr@paperweekly.site
• 来稿请备注即时联系方式（微信），以便我们在稿件选用的第一时间联系作者
• 您也可以直接添加小编微信（pwbot02）快速投稿，备注：姓名-投稿
△长按添加PaperWeekly小编
🔍
现在，在「知乎」也能找到我们了
进入知乎首页搜索「PaperWeekly」
点击「关注」订阅我们的专栏吧